import os, sys
sys.path.append('..')
import pandas as pd
import numpy as np
from tqdm import tqdm
import platform
import cv2
import torch
import segmentation_models_pytorch as smp
from torch.utils.data import DataLoader, Dataset
from PIL import Image
from albumentations.pytorch import ToTensor
from albumentations import(
    Resize, Normalize, Compose
)

batch_size = 12
MASK_TH = 0.5
img_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_test_data/Images'
dst_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_test_data/seg_pred_mask'
pth_path = '/opt/data/private/project/adc_segmentation/ckpts/50UD_05300/erxian/Unet_se_resnext50_32x4d_latest_epoch.pth'

df = pd.DataFrame()
codes = []
imgs = []
for code in os.listdir(img_path):
    tmp = os.listdir(os.path.join(img_path, code))
    tmp2 = []
    for x in tmp:
        if not os.path.exists(os.path.join(dst_path, code, x[:-3]+'png')):
            tmp2.append(x)
    codes.extend([code] * len(tmp2))
    imgs.extend(tmp2[:])

df['root_path'] = [img_path] * len(codes)
df['code'] = codes
df['image'] = imgs


def resize_transform_basic(img_size = (256, 256), mean=[0,0,0], std=[1,1,1]):
    '''
    Transformations only include (Resize and normalize)
    '''
    # img_size = (height, width)
    # crop_size = (height, width)
    if isinstance(img_size, int):
        img_size = (img_size, img_size)
    transforms_list = []
    transforms_list.extend([
        Resize(img_size[0], img_size[1], interpolation=Image.BILINEAR),
        Normalize(mean=mean, std=std, p=1),
        ToTensor(),
        ])
    transforms = Compose(transforms_list)
    return transforms

class T2_Seg_Dataset(Dataset):
    def __init__(self, dataframe, transform=None):
        self.df = dataframe
        if transform is None:
            self.seg_transforms = resize_transform_basic()
        else:
            self.seg_transforms = transform

    def __len__(self):
        return self.df.shape[0]

    def __getitem__(self, idx):
        root_path = self.df["root_path"].iloc[idx]
        code = self.df["code"].iloc[idx]
        img_name = self.df["image"].iloc[idx]
        try:
            image = cv2.imread(os.path.join(root_path, code, img_name))
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        except IOError:
            print(code, img_name)
            raise ValueError("The image path is incorrect, please check the image path.")
        aug = self.seg_transforms(image=image)
        return aug['image'], code, img_name


dataset = T2_Seg_Dataset(df, resize_transform_basic(img_size=(640, 480),
                                                    mean=[0.485, 0.456, 0.406],
                                                    std=[0.229, 0.224, 0.225]))
kwargs = {'num_workers': 4, 'pin_memory': True} if (torch.cuda.is_available() and platform.system() == 'Linux') else {}
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, **kwargs)
model = smp.Unet('se_resnext50_32x4d', encoder_weights=None, classes=1, activation=None)
# params = torch.load('pre_trained_model/model.pth', map_location=torch.device('cpu'))
params = torch.load(pth_path, map_location=torch.device('cpu'))
model.load_state_dict(params["state_dict"])
model.eval()
if torch.cuda.is_available():
    model.cuda()

for itr, (images, codes, img_names) in enumerate(tqdm(loader)):
    images = images.cuda()
    with torch.no_grad():
        outputs = model(images).float()
        result = torch.sigmoid(outputs)
        probability = result.detach().cpu().numpy()
        probability = np.squeeze(probability)
    for i in range(len(codes)):
        mask = cv2.threshold(probability[i], MASK_TH, 1, cv2.THRESH_BINARY)[1]
        os.makedirs(os.path.join(dst_path, codes[i]), exist_ok=True)
        cv2.imwrite(os.path.join(dst_path, codes[i], img_names[i][:-3]+'png'), mask)
